#include "ego_planner_adapter/bspline_interpolator.hpp"
#include <algorithm>
#include <cmath>
#include <iostream>

// Include ego-planner messages if available
#ifdef TRAJ_UTILS_FOUND
#include "traj_utils/msg/bspline.hpp"
#include "geometry_msgs/msg/point.hpp"
#endif

namespace ego_planner_adapter {

BSplineInterpolator::BSplineInterpolator()
  : yaw_dt_(0.0),
    start_time_(0.0),
    end_time_(0.0),
    order_(3),
    is_valid_(false),
    num_control_points_(0),
    num_knots_(0),
    trajectory_duration_(0.0) {
}

#ifdef TRAJ_UTILS_FOUND
bool BSplineInterpolator::setTrajectory(const std::shared_ptr<const traj_utils::msg::Bspline>& bspline_msg) {
    clear();
    
    if (!bspline_msg) {
        std::cerr << "BSplineInterpolator: Null trajectory message" << std::endl;
        return false;
    }
    
    // Extract trajectory parameters
    order_ = bspline_msg->order;
    
    // Convert control points
    control_points_.clear();
    control_points_.reserve(bspline_msg->pos_pts.size());
    
    for (const auto& point : bspline_msg->pos_pts) {
        control_points_.emplace_back(point.x, point.y, point.z);
    }
    
    // Copy knot sequence
    knots_ = bspline_msg->knots;
    
    // Copy yaw points
    yaw_points_ = bspline_msg->yaw_pts;
    yaw_dt_ = bspline_msg->yaw_dt;
    
    // Extract timing information
    auto start_time_msg = bspline_msg->start_time;
    start_time_ = start_time_msg.sec + start_time_msg.nanosec * 1e-9;
    
    // Validate trajectory data
    num_control_points_ = control_points_.size();
    num_knots_ = knots_.size();
    
    if (num_control_points_ < order_ + 1) {
        std::cerr << "BSplineInterpolator: Insufficient control points (" 
                  << num_control_points_ << " < " << order_ + 1 << ")" << std::endl;
        return false;
    }
    
    if (num_knots_ != num_control_points_ + order_ + 1) {
        std::cerr << "BSplineInterpolator: Incorrect knot sequence length ("
                  << num_knots_ << " != " << num_control_points_ + order_ + 1 << ")" << std::endl;
        return false;
    }
    
    // Compute trajectory duration from knot sequence
    if (knots_.size() >= 2) {
        trajectory_duration_ = knots_.back() - knots_[order_];
        end_time_ = start_time_ + trajectory_duration_;
    } else {
        std::cerr << "BSplineInterpolator: Invalid knot sequence" << std::endl;
        return false;
    }
    
    // Validate knot sequence is non-decreasing
    for (size_t i = 1; i < knots_.size(); ++i) {
        if (knots_[i] < knots_[i-1]) {
            std::cerr << "BSplineInterpolator: Non-monotonic knot sequence at index " << i << std::endl;
            return false;
        }
    }
    
    is_valid_ = true;
    
    std::cout << "BSplineInterpolator: Trajectory set successfully\n"
              << "  Control points: " << num_control_points_ << "\n"
              << "  Knots: " << num_knots_ << "\n"
              << "  Order: " << order_ << "\n"
              << "  Duration: " << trajectory_duration_ << " seconds\n"
              << "  Start time: " << start_time_ << "\n"
              << "  End time: " << end_time_ << std::endl;
    
    return true;
}
#else
bool BSplineInterpolator::setTrajectory(const std::shared_ptr<const traj_utils::msg::Bspline>& /*bspline_msg*/) {
    std::cerr << "BSplineInterpolator: traj_utils not available" << std::endl;
    return false;
}
#endif

bool BSplineInterpolator::getPositionAtTime(double time, Eigen::Vector3d& position) {
    if (!is_valid_) {
        return false;
    }
    
    // Convert time to parameter
    double t = timeToParameter(time);
    
    // Clamp parameter to valid range
    t = std::clamp(t, 0.0, 1.0);
    
    // Evaluate B-spline
    position = evaluateBSpline(t);
    
    return true;
}

bool BSplineInterpolator::getVelocityAtTime(double time, Eigen::Vector3d& velocity) {
    if (!is_valid_) {
        return false;
    }
    
    // Convert time to parameter
    double t = timeToParameter(time);
    
    // Clamp parameter to valid range
    t = std::clamp(t, 0.0, 1.0);
    
    // Evaluate B-spline derivative
    velocity = evaluateBSplineDerivative(t);
    
    // Scale by parameter derivative (dt/ds)
    if (trajectory_duration_ > EPSILON) {
        velocity /= trajectory_duration_;
    }
    
    return true;
}

bool BSplineInterpolator::getYawAtTime(double time, double& yaw) {
    if (!is_valid_) {
        return false;
    }
    
    if (yaw_points_.empty() || yaw_dt_ <= EPSILON) {
        yaw = 0.0;
        return true;
    }
    
    // Convert time to yaw index
    double relative_time = time - start_time_;
    int yaw_index = static_cast<int>(relative_time / yaw_dt_);
    
    // Clamp to valid range
    if (yaw_index < 0) {
        yaw = yaw_points_[0];
    } else if (yaw_index >= static_cast<int>(yaw_points_.size())) {
        yaw = yaw_points_.back();
    } else {
        // Linear interpolation between yaw points
        double alpha = (relative_time / yaw_dt_) - yaw_index;
        
        if (yaw_index + 1 < static_cast<int>(yaw_points_.size())) {
            yaw = yaw_points_[yaw_index] * (1.0 - alpha) + yaw_points_[yaw_index + 1] * alpha;
        } else {
            yaw = yaw_points_[yaw_index];
        }
    }
    
    return true;
}

void BSplineInterpolator::clear() {
    control_points_.clear();
    knots_.clear();
    yaw_points_.clear();
    yaw_dt_ = 0.0;
    start_time_ = 0.0;
    end_time_ = 0.0;
    order_ = 3;
    is_valid_ = false;
    num_control_points_ = 0;
    num_knots_ = 0;
    trajectory_duration_ = 0.0;
}

Eigen::Vector3d BSplineInterpolator::evaluateBSpline(double t) {
    if (!is_valid_ || control_points_.empty()) {
        return Eigen::Vector3d::Zero();
    }
    
    // Convert normalized parameter t to knot parameter
    double u = knots_[order_] + t * (knots_[num_control_points_] - knots_[order_]);
    
    Eigen::Vector3d result = Eigen::Vector3d::Zero();
    
    // Evaluate using de Boor's algorithm
    for (int i = 0; i < num_control_points_; ++i) {
        double basis = basisFunction(i, order_, u);
        result += basis * control_points_[i];
    }
    
    return result;
}

Eigen::Vector3d BSplineInterpolator::evaluateBSplineDerivative(double t) {
    if (!is_valid_ || control_points_.empty() || order_ < 1) {
        return Eigen::Vector3d::Zero();
    }
    
    // Convert normalized parameter t to knot parameter
    double u = knots_[order_] + t * (knots_[num_control_points_] - knots_[order_]);
    
    Eigen::Vector3d result = Eigen::Vector3d::Zero();
    
    // Evaluate derivative using derivative of basis functions
    for (int i = 0; i < num_control_points_; ++i) {
        double basis_deriv = basisFunctionDerivative(i, order_, u);
        result += basis_deriv * control_points_[i];
    }
    
    return result;
}

double BSplineInterpolator::timeToParameter(double time) {
    if (trajectory_duration_ <= EPSILON) {
        return 0.0;
    }
    
    double relative_time = time - start_time_;
    return relative_time / trajectory_duration_;
}

double BSplineInterpolator::basisFunction(int i, int k, double t) {
    if (k == 0) {
        // Base case: degree 0 basis function
        return (t >= knots_[i] && t < knots_[i + 1]) ? 1.0 : 0.0;
    }
    
    double result = 0.0;
    
    // Left term
    double denom1 = knots_[i + k] - knots_[i];
    if (std::abs(denom1) > EPSILON) {
        result += (t - knots_[i]) / denom1 * basisFunction(i, k - 1, t);
    }
    
    // Right term
    double denom2 = knots_[i + k + 1] - knots_[i + 1];
    if (std::abs(denom2) > EPSILON) {
        result += (knots_[i + k + 1] - t) / denom2 * basisFunction(i + 1, k - 1, t);
    }
    
    return result;
}

double BSplineInterpolator::basisFunctionDerivative(int i, int k, double t) {
    if (k == 0) {
        return 0.0; // Derivative of step function is zero (except at discontinuities)
    }
    
    double result = 0.0;
    
    // Left term
    double denom1 = knots_[i + k] - knots_[i];
    if (std::abs(denom1) > EPSILON) {
        result += k / denom1 * basisFunction(i, k - 1, t);
    }
    
    // Right term
    double denom2 = knots_[i + k + 1] - knots_[i + 1];
    if (std::abs(denom2) > EPSILON) {
        result -= k / denom2 * basisFunction(i + 1, k - 1, t);
    }
    
    return result;
}

} // namespace ego_planner_adapter